"""
演示数据生成器
生成用于演示的慢性病风险预测数据
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import os

class DemoDataGenerator:
    """演示数据生成器"""
    
    def __init__(self, seed=42):
        np.random.seed(seed)
        
    def generate_patient_data(self, n_patients=1000):
        """生成患者基础数据"""
        print(f"正在生成 {n_patients} 个患者的演示数据...")
        
        # 基础信息
        ages = np.random.normal(50, 15, n_patients).astype(int)
        ages = np.clip(ages, 18, 100)  # 限制年龄范围
        
        genders = np.random.choice([0, 1], n_patients, p=[0.5, 0.5])  # 0: 女, 1: 男
        
        # 身高体重（基于性别调整）
        heights = np.where(genders == 1, 
                          np.random.normal(175, 8, n_patients),  # 男性
                          np.random.normal(162, 6, n_patients))  # 女性
        heights = np.clip(heights, 140, 200)
        
        weights = np.where(genders == 1,
                          np.random.normal(75, 12, n_patients),  # 男性
                          np.random.normal(60, 10, n_patients))  # 女性
        weights = np.clip(weights, 40, 150)
        
        # 计算BMI
        bmis = weights / ((heights / 100) ** 2)
        
        # 生理指标（基于年龄和BMI调整）
        # 血压
        systolic_bp = np.random.normal(120, 15, n_patients)
        systolic_bp += (ages - 50) * 0.3  # 年龄影响
        systolic_bp += (bmis - 23) * 0.8  # BMI影响
        systolic_bp = np.clip(systolic_bp, 90, 200)
        
        diastolic_bp = systolic_bp * 0.67 + np.random.normal(0, 5, n_patients)
        diastolic_bp = np.clip(diastolic_bp, 50, 120)
        
        # 血糖
        blood_glucose = np.random.normal(100, 15, n_patients)
        blood_glucose += (ages - 50) * 0.2  # 年龄影响
        blood_glucose += (bmis - 23) * 1.2  # BMI影响
        blood_glucose = np.clip(blood_glucose, 70, 200)
        
        # 胆固醇
        cholesterol = np.random.normal(200, 30, n_patients)
        cholesterol += (ages - 50) * 0.5
        cholesterol = np.clip(cholesterol, 120, 350)
        
        # 心率
        heart_rate = np.random.normal(72, 10, n_patients)
        heart_rate = np.clip(heart_rate, 50, 120)
        
        # 生活方式因素
        smoking = np.random.choice([0, 1], n_patients, p=[0.7, 0.3])
        drinking = np.random.choice([0, 1], n_patients, p=[0.6, 0.4])
        family_history = np.random.choice([0, 1], n_patients, p=[0.8, 0.2])
        
        # 运动频率 (0: 从不, 1: 偶尔, 2: 经常, 3: 每天)
        exercise_frequency = np.random.choice([0, 1, 2, 3], n_patients, p=[0.2, 0.3, 0.3, 0.2])
        
        # 饮食偏好 (0: 清淡, 1: 正常, 2: 重口味)
        diet_preference = np.random.choice([0, 1, 2], n_patients, p=[0.3, 0.5, 0.2])
        
        # 生成标签（基于风险因素）
        # 高血压风险
        ht_risk_score = (
            (ages - 30) * 0.01 +
            (bmis - 20) * 0.02 +
            (systolic_bp - 120) * 0.005 +
            smoking * 0.15 +
            family_history * 0.1 +
            (3 - exercise_frequency) * 0.05 +
            np.random.normal(0, 0.1, n_patients)
        )
        hypertension_risk = (ht_risk_score > 0.3).astype(int)
        
        # 糖尿病风险
        dm_risk_score = (
            (ages - 30) * 0.008 +
            (bmis - 20) * 0.025 +
            (blood_glucose - 100) * 0.003 +
            smoking * 0.1 +
            family_history * 0.12 +
            (3 - exercise_frequency) * 0.08 +
            np.random.normal(0, 0.1, n_patients)
        )
        diabetes_risk = (dm_risk_score > 0.25).astype(int)
        
        # 创建DataFrame
        data = pd.DataFrame({
            'patient_id': [f'P{i+1:04d}' for i in range(n_patients)],
            'age': ages,
            'gender': genders,
            'height': heights.astype(int),
            'weight': weights.astype(int),
            'bmi': bmis.round(1),
            'systolic_bp': systolic_bp.astype(int),
            'diastolic_bp': diastolic_bp.astype(int),
            'blood_glucose': blood_glucose.astype(int),
            'cholesterol': cholesterol.astype(int),
            'heart_rate': heart_rate.astype(int),
            'smoking': smoking,
            'drinking': drinking,
            'family_history': family_history,
            'exercise_frequency': exercise_frequency,
            'diet_preference': diet_preference,
            'hypertension_risk': hypertension_risk,
            'diabetes_risk': diabetes_risk
        })
        
        return data
    
    def generate_timeseries_data(self, patient_ids, n_days=90):
        """生成时序生理指标数据"""
        print(f"正在生成 {len(patient_ids)} 个患者的 {n_days} 天时序数据...")
        
        timeseries_data = []
        
        for patient_id in patient_ids:
            # 获取患者基础信息（这里简化处理）
            base_systolic = np.random.normal(120, 10)
            base_diastolic = base_systolic * 0.67
            base_glucose = np.random.normal(100, 8)
            base_heart_rate = np.random.normal(72, 5)
            
            # 生成日期序列
            start_date = datetime.now() - timedelta(days=n_days)
            dates = [start_date + timedelta(days=i) for i in range(n_days)]
            
            # 生成带趋势的时序数据
            for i, date in enumerate(dates):
                # 添加趋势和随机波动
                trend_factor = i * 0.01  # 轻微上升趋势
                noise_factor = np.random.normal(0, 0.05)
                
                systolic = base_systolic + trend_factor + noise_factor * 5
                diastolic = base_diastolic + trend_factor * 0.7 + noise_factor * 3
                glucose = base_glucose + trend_factor * 0.5 + noise_factor * 4
                heart_rate = base_heart_rate + noise_factor * 3
                
                timeseries_data.append({
                    'patient_id': patient_id,
                    'date': date.strftime('%Y-%m-%d'),
                    'systolic_bp': max(80, min(180, systolic)),
                    'diastolic_bp': max(50, min(110, diastolic)),
                    'blood_glucose': max(70, min(150, glucose)),
                    'heart_rate': max(50, min(100, heart_rate))
                })
        
        return pd.DataFrame(timeseries_data)
    
    def save_demo_data(self, output_dir='demo_data'):
        """保存演示数据"""
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成患者数据
        patient_data = self.generate_patient_data(1000)
        patient_data.to_csv(f'{output_dir}/patient_data.csv', index=False, encoding='utf-8-sig')
        print(f"患者数据已保存到: {output_dir}/patient_data.csv")
        
        # 生成时序数据（选择前100个患者）
        selected_patients = patient_data['patient_id'].head(100).tolist()
        timeseries_data = self.generate_timeseries_data(selected_patients, 90)
        timeseries_data.to_csv(f'{output_dir}/timeseries_data.csv', index=False, encoding='utf-8-sig')
        print(f"时序数据已保存到: {output_dir}/timeseries_data.csv")
        
        # 生成统计报告
        self.generate_data_report(patient_data, f'{output_dir}/data_report.txt')
        
        return patient_data, timeseries_data
    
    def generate_data_report(self, data, output_file):
        """生成数据统计报告"""
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("慢性病风险预测演示数据统计报告\n")
            f.write("=" * 50 + "\n\n")
            
            f.write(f"数据生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"总患者数: {len(data)}\n\n")
            
            f.write("基础统计信息:\n")
            f.write("-" * 30 + "\n")
            f.write(f"平均年龄: {data['age'].mean():.1f} 岁\n")
            f.write(f"年龄范围: {data['age'].min()}-{data['age'].max()} 岁\n")
            f.write(f"平均BMI: {data['bmi'].mean():.1f}\n")
            f.write(f"平均收缩压: {data['systolic_bp'].mean():.1f} mmHg\n")
            f.write(f"平均舒张压: {data['diastolic_bp'].mean():.1f} mmHg\n")
            f.write(f"平均血糖: {data['blood_glucose'].mean():.1f} mg/dL\n\n")
            
            f.write("风险分布:\n")
            f.write("-" * 30 + "\n")
            f.write(f"高血压风险患者: {data['hypertension_risk'].sum()} 人 ({data['hypertension_risk'].mean()*100:.1f}%)\n")
            f.write(f"糖尿病风险患者: {data['diabetes_risk'].sum()} 人 ({data['diabetes_risk'].mean()*100:.1f}%)\n\n")
            
            f.write("生活方式统计:\n")
            f.write("-" * 30 + "\n")
            f.write(f"吸烟者: {data['smoking'].sum()} 人 ({data['smoking'].mean()*100:.1f}%)\n")
            f.write(f"饮酒者: {data['drinking'].sum()} 人 ({data['drinking'].mean()*100:.1f}%)\n")
            f.write(f"有家族史: {data['family_history'].sum()} 人 ({data['family_history'].mean()*100:.1f}%)\n")
            
            exercise_dist = data['exercise_frequency'].value_counts().sort_index()
            f.write(f"运动频率分布:\n")
            for freq, count in exercise_dist.items():
                freq_name = ['从不', '偶尔', '经常', '每天'][freq]
                f.write(f"  {freq_name}: {count} 人 ({count/len(data)*100:.1f}%)\n")
        
        print(f"数据统计报告已保存到: {output_file}")

def main():
    """主函数"""
    print("慢性病风险预测系统 - 演示数据生成器")
    print("=" * 50)
    
    generator = DemoDataGenerator()
    patient_data, timeseries_data = generator.save_demo_data()
    
    print("\n演示数据生成完成！")
    print(f"患者数据: {len(patient_data)} 条记录")
    print(f"时序数据: {len(timeseries_data)} 条记录")
    print("\n数据文件位置:")
    print("- demo_data/patient_data.csv")
    print("- demo_data/timeseries_data.csv")
    print("- demo_data/data_report.txt")

if __name__ == "__main__":
    main()



